package com.pai4j.aigc.llm;

import com.pai4j.aigc.service.BotService;
import com.pai4j.aigc.service.transfunction.AbstractTransFuncService;
import com.pai4j.aigc.service.transfunction.TransFunctionFactory;
import com.pai4j.aigc.llm.entity.LLMModelEntity;
import com.pai4j.aigc.llm.service.LLMModelService;
import com.pai4j.aigc.llm.service.LLMUsageService;
import com.pai4j.aigc.llm.service.MetricsService;
import com.pai4j.aigc.llm.service.PricingService;
import com.pai4j.aigc.llm.service.TokenEstimator;
import com.pai4j.common.helper.SessionHelper;
import com.pai4j.common.service.FreeMarkEngineService;
import com.pai4j.common.util.DateUtils;
import com.pai4j.common.util.JsonExtractorFromMarkdown;
import com.pai4j.common.util.JsonUtil;
import com.pai4j.domain.dto.bot.BotDTO;
import com.pai4j.domain.dto.bot.BotSchemaDTO;
import com.pai4j.domain.dto.bot.LLMParameterDTO;
import com.pai4j.domain.dto.bot.PluginParameterTransRuleDTO;
import com.pai4j.domain.vo.llm.ChatCompletionMessage;
import com.pai4j.domain.dto.llm.ChatMessageRole;
import com.pai4j.domain.vo.llm.ChatCompletionChoice;
import com.pai4j.domain.vo.llm.ChatCompletionResponse;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import com.pai4j.common.enums.PromptTemplateEnum;
import java.util.*;


/**
 * 用户意图识别后置抽参
 */
@Slf4j
@Service
public class AiParameterService {

    @Value("${llm.config.model}")
    private String model;

    @Resource
    private FreeMarkEngineService freeMarkEngineService;

    @Resource
    private BotService botService;

    @Resource
    private LLMUsageService usageService;

    @Resource
    private TokenEstimator tokenEstimator;

    @Resource
    private PricingService pricingService;

    @Resource
    private MetricsService metricsService;

    @Resource
    private LLMModelService modelService;

    /**
     * 用户输入意图识别
     *
     * @return
     */
    public LLMParameterDTO getParameter(String currentUserId, String input, String botCode) throws Exception {
        // 查询bot详情
        BotDTO botDTO = botService.getBotPluginFromCache(botCode);
        // 插件参数抽取
        LLMParameterDTO parameterDTO = this.getParameterByLLM(botDTO, input);
        if (parameterDTO == null) {
            return null;
        }
        /**
         * 根据schema配置对抽取参数value进行转换，转换成插件API能够认识的参数
         */
        this.transParameter(currentUserId, botDTO.getSchema(), parameterDTO);
        // 设置当前bot的插件接口配置
        parameterDTO.setPlugin(botDTO.getPlugin());
        return parameterDTO;
    }


    private LLMParameterDTO getParameterByLLM(BotDTO botDTO, String input) throws Exception {
        String todayDate = DateUtils.getCurrentFormatDate(DateUtils.YYYY_MM_DD);
        Long todayStartTime = DateUtils.getStartOfDay(new Date()).getTime();
        Map<String, Object> vars = new HashMap<>();
        vars.put("schema", botDTO.getSchema());
        vars.put("rule", botDTO.getRule());
        vars.put("outputExample", botDTO.getOutputExample());
        vars.put("input", input);
        vars.put("todayDate", todayDate);
        vars.put("todayStartTime", todayStartTime);
        String prompt = freeMarkEngineService.getContentByTemplate(PromptTemplateEnum.AI_BOT_PARAMETER.getPath(), vars);
        // 调用大模型，意图识别
        final List<ChatCompletionMessage> messageList = Arrays.asList(
                new ChatCompletionMessage(ChatMessageRole.USER.value(), prompt));
        AbstractLLMChatService abstractLLMChatService = LLMServiceFactory.getLLMService(model);
        // resolve model
        java.util.Optional<LLMModelEntity> modelEntityOpt = modelService.findByGlobalCode(model);
        LLMModelEntity modelEntity = modelEntityOpt.orElse(null);
        String providerKey = modelEntity != null ? modelEntity.getProvider() : model;
        Long modelId = modelEntity != null ? modelEntity.getId() : -1L;
        String requestId = java.util.UUID.randomUUID().toString();
        long start = System.currentTimeMillis();
        try {
            ChatCompletionResponse chatCompletionResponse = abstractLLMChatService.chat(messageList);
            ChatCompletionChoice chatCompletionChoice = chatCompletionResponse.getChoices().get(0);
            String functionMdJson = chatCompletionChoice.getMessage().getContent();

            long latency = System.currentTimeMillis() - start;
            int pt = tokenEstimator.estimateTokens(prompt);
            int ct = tokenEstimator.estimateTokens(functionMdJson);
            int total = pt + ct;
            Integer costCents = pricingService.calcCostCents(modelEntity, pt, ct);
            usageService.recordSuccess(null, requestId, null, modelId, providerKey,
                    modelEntity != null ? modelEntity.getCode() : providerKey,
                    pt, ct, total, latency, costCents);
            metricsService.recordSuccess(providerKey, null, modelEntity != null ? modelEntity.getCode() : providerKey, latency);

            String functionJson = JsonExtractorFromMarkdown.getJson(functionMdJson);
            return JsonUtil.fromJson(functionJson, LLMParameterDTO.class);
        } catch (Exception e) {
            long latency = System.currentTimeMillis() - start;
            usageService.recordFailure(null, requestId, null, modelId, providerKey,
                    modelEntity != null ? modelEntity.getCode() : providerKey,
                    "LLM_ERROR", e.getMessage(), latency);
            metricsService.recordFailure(providerKey, null, modelEntity != null ? modelEntity.getCode() : providerKey, latency);
            throw e;
        }
    }

    private void transParameter(String currentUserId, String schema, LLMParameterDTO parameterDTO) {
        if (StringUtils.isBlank(schema)) {
            return;
        }
        BotSchemaDTO botSchemaDTO = JsonUtil.fromJson(schema, BotSchemaDTO.class);
        List<PluginParameterTransRuleDTO> transRuleDTOS = botSchemaDTO.getTransRules();
        if (CollectionUtils.isEmpty(transRuleDTOS)) {
            return;
        }
        // 当前登录用户添加到当前现成副本，后续主线程任务需要时直接获取
        SessionHelper.setContextUserId(currentUserId);
        for (PluginParameterTransRuleDTO transRule : transRuleDTOS) {
            String parameterName = transRule.getParameter();
            String transParameterName = transRule.getTransParameter();
            String rule = transRule.getRule();
            // 获取当前需要转换的参数名对应的参数值
            Object parameterValue = parameterDTO.getParameterValue(parameterName);
            // 走抽象工厂模式实现参数trans
            AbstractTransFuncService abstractTransFuncService = TransFunctionFactory.getService(rule);
            Object transParameterValue = abstractTransFuncService.trans(parameterValue);
            // 重置转换后的参数名及参数值，用于api请求
            parameterDTO.resetParameterValue(parameterName, transParameterName, String.valueOf(transParameterValue));
        }
    }

}
